package edu.zju.gis.dbfg.model.statistic;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.base.Constant;
import edu.zju.gis.dbfg.model.datasource.DataReaderRDD;
import edu.zju.gis.dbfg.model.util.CrsPick;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.StatisticIndex;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.api.java.function.Function2;
import org.apache.spark.api.java.function.PairFunction;
import org.apache.spark.broadcast.Broadcast;
import org.geotools.referencing.CRS;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.File;
import java.io.FileWriter;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


public class PacStatisticIndex {

    private final static Logger logger = LoggerFactory.getLogger(PacStatisticIndex.class);

    public static void main(String[] args) throws Exception{
        //Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
//        conf.setMaster("local[4]");
//        conf.setAppName("PacStatisticIndex");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        //Setup model arguments
        logger.info("Setup arguments");
        PacStatisticIndexArgs mArgs = PacStatisticIndexArgs.initArgs(args);
        if (mArgs == null){
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(PacStatisticIndex.class, "main()", msg, args);
        }

        CoordinateReferenceSystem crs1 = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs1()));
        CoordinateReferenceSystem targetCRS = Constant.DEFAULT_CRS();
        MathTransform transform1 = null;
        if(!crs1.equals(targetCRS)){
            transform1 = CRS.findMathTransform(crs1, targetCRS);
        }
        Broadcast<PacStatisticIndexArgs> mArgsBroad = jsc.broadcast(mArgs);

        String input1 = mArgs.getInput();
        String type1 = mArgs.getType();
        String cc1 = mArgs.getCcIndex();
        String pac1 = mArgs.getPacIndex();
        String area1 = mArgs.getAreaIndex();
        List<String> indecies = Arrays.asList(StatisticIndex.getIndecies(mArgs.getIndexCategory()));
        Map<String, Object> attrs1 = new HashMap<>();
        attrs1.put("CC",cc1);
        attrs1.put("PAC",pac1);
        if(!area1.equals("-1")){
            attrs1.put("AREA", area1);
        }
        JavaRDD<GeoItem> gi1 = null;
        List<String> filters = Arrays.asList(mArgs.getFilters().split(","));
        if(type1.equals("ES")){
            gi1 = DataReaderRDD.readToRDD2(type1, jsc, input1, attrs1,filters, mArgs.getFilterIndex(), transform1);
        }else {
            gi1 = DataReaderRDD.readToRDD(type1, jsc, input1, null, attrs1, transform1);
        }

        JavaPairRDD<String,GeoItem> resultRDD = gi1.mapToPair(new PairFunction<GeoItem, String, GeoItem>() {
            @Override
            public Tuple2<String, GeoItem> call(GeoItem gi1) throws Exception {
                String cc = String.valueOf(gi1.getProp().get("CC"));
                if((null == indecies) || ((null != indecies)&&(indecies.contains(cc))))
                {
                    String pac = String.valueOf(gi1.getProp().get("PAC"));
                    String id = String.format("%s##%s",pac,cc);
                    return new Tuple2<>(id, gi1);
                }
                return new Tuple2<>("EMPTY",null);
            }
        }).filter(new Function<Tuple2<String, GeoItem>, Boolean>() {
            @Override
            public Boolean call(Tuple2<String, GeoItem> input) throws Exception {
                return !input._1.equals("EMPTY");
            }
        });

        JavaPairRDD<String, GeoItem> resultcc = resultRDD.reduceByKey(new Function2<GeoItem, GeoItem, GeoItem>() {
            @Override
            public GeoItem call(GeoItem gi1, GeoItem gi2) throws Exception {
                Double area = (Double)gi1.getProp().get("AREA")+(Double)gi2.getProp().get("AREA");
                GeoItem result = new GeoItem();
                result.setFid(gi1.getFid());
                result.setProp(gi1.getProp());
                result.getProp().put("AREA",area);
                return result;
            }
        });

        Map<String, GeoItem> result = resultcc.collectAsMap();
        File file = new File(mArgs.getOutput());
        if(file.exists()){
            file.delete();
        }
        FileWriter writer = new FileWriter(file);
        if(area1.equals("-1")){
            area1 = "area";
        }
        writer.write(String.format("%s\t%s\t%s\n", pac1, cc1, area1));
        for(String key: result.keySet()){
            String[] fields = key.split("##");
            String pac = fields[0];
            String cc = fields[1];
            GeoItem gi = result.get(key);
            Double area = (Double) gi.getProp().get("AREA");
            String s = String.format("%s\t%s\t%.9f", pac, cc, area);
            writer.write(s + "\n");
        }
        writer.flush();
        writer.close();

    }
}
